# -*- coding: utf-8 -*-

import matplotlib.pyplot as plt
from matplotlib import cm
import pickle
import pandas as pd


def plot_HRdiagram(BV, Vmag):
    '''
    画出 毕罗图，其中颜色根据色指数
    '''
    
    fig, ax = plt.subplots()
    cmap = cm.hot
    ax.scatter(Vmag, BV, c=BV, cmap=cmap)
    ax.set_xlabel(r'B-V', fontsize=15)
    ax.set_ylabel(r'Vmag', fontsize=15)
    ax.set_title('H-R diagram')
    ax.grid(True)
    ax.set_facecolor((0, 0, 0))
    fig.colorbar(cm.ScalarMappable(cmap=cmap))
    plt.show()

if __name__ == '__main__':
    bi_clu_label = pickle.load(open(r'../附件/bi_star_cluster_label.pkl', 'rb'))
    cluster_res = pickle.load(open(r'../附件/clustering_result.pkl', 'rb'))
    path = r'../附件/星表数据.txt'
    # 读取 txt 文件
    df = pd.read_csv(path, delimiter=' +')
    df['label'] = cluster_res
    # 找出 【20,22】 的点
    Plx_20_22 = df.loc[(20<=df['Plx']) & (df['Plx']<=22)].index
    df = df.iloc[Plx_20_22]
    # 提取毕星团
    bi_stars = df.loc[df['label']==bi_clu_label]

    x = bi_stars['Vmag']
    y = bi_stars['B-V']
    plot_HRdiagram(x, y)
